#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
快速设置环境变量脚本（交互版）
- 支持交互选择与输入
- 写出 transfer_optimization.env
- 提示前端点击“重载配置”或调用 /api/reload_config
"""

import os
from datetime import datetime

DEFAULTS = {
    'TRANSFER_CHUNK_SIZE': 16 * 1024 * 1024,   # 16MB
    'TRANSFER_USE_COMPRESSION': 'true',
    'TRANSFER_COMPRESSION_LEVEL': 6,
    'TRANSFER_USE_PARALLEL': 'true',
    'TRANSFER_PARALLEL_CHUNKS': 3,
    'TRANSFER_USE_HTTP': 'false',
    'TRANSFER_USE_RSYNC': 'false',
    'DOCKER_TIMEOUT': 300,
    'DOCKER_PULL_TIMEOUT': 600,
    'DOCKER_SAVE_TIMEOUT': 900,
}

YES = {'y','yes','true','1'}
NO = {'n','no','false','0'}


def prompt_bool(name:str, default_str:str) -> str:
    default_disp = 'Y/n' if default_str == 'true' else 'y/N'
    while True:
        s = input(f"{name} ({default_disp}): ").strip().lower()
        if not s:
            return default_str
        if s in YES:
            return 'true'
        if s in NO:
            return 'false'
        print('请输入 y 或 n')


def prompt_int(name:str, default_val:int, min_val:int=None, max_val:int=None) -> int:
    while True:
        s = input(f"{name} (默认 {default_val}): ").strip()
        if not s:
            return default_val
        try:
            v = int(s)
            if min_val is not None and v < min_val:
                print(f"必须 >= {min_val}")
                continue
            if max_val is not None and v > max_val:
                print(f"必须 <= {max_val}")
                continue
            return v
        except ValueError:
            print('请输入整数')


def prompt_size_mb(name:str, default_mb:int, min_mb:int=1, max_mb:int=128) -> int:
    v = prompt_int(name + ' (MB)', default_mb, min_mb, max_mb)
    return v * 1024 * 1024


def main():
    print('=== Docker 镜像传输优化 环境配置 ===')
    chunk_size = prompt_size_mb('切片大小', DEFAULTS['TRANSFER_CHUNK_SIZE'] // (1024*1024))
    use_compress = prompt_bool('启用压缩', DEFAULTS['TRANSFER_USE_COMPRESSION'])
    level = prompt_int('压缩级别(1-9)', DEFAULTS['TRANSFER_COMPRESSION_LEVEL'], 1, 9)
    use_parallel = prompt_bool('启用并行传输', DEFAULTS['TRANSFER_USE_PARALLEL'])
    parallel_chunks = prompt_int('并行数量(1-5)', DEFAULTS['TRANSFER_PARALLEL_CHUNKS'], 1, 5)
    use_http = prompt_bool('使用HTTP传输(仅内网建议)', DEFAULTS['TRANSFER_USE_HTTP'])
    use_rsync = prompt_bool('使用rsync传输(需目标机安装rsync)', DEFAULTS['TRANSFER_USE_RSYNC'])

    docker_timeout = prompt_int('DOCKER_TIMEOUT(秒)', DEFAULTS['DOCKER_TIMEOUT'])
    pull_timeout = prompt_int('DOCKER_PULL_TIMEOUT(秒)', DEFAULTS['DOCKER_PULL_TIMEOUT'])
    save_timeout = prompt_int('DOCKER_SAVE_TIMEOUT(秒)', DEFAULTS['DOCKER_SAVE_TIMEOUT'])

    env_text = f"""# Docker镜像传输优化配置
# 生成时间: {datetime.now():%Y-%m-%d %H:%M:%S}
TRANSFER_CHUNK_SIZE={chunk_size}
TRANSFER_USE_COMPRESSION={use_compress}
TRANSFER_COMPRESSION_LEVEL={level}
TRANSFER_USE_PARALLEL={use_parallel}
TRANSFER_PARALLEL_CHUNKS={parallel_chunks}
TRANSFER_USE_HTTP={use_http}
TRANSFER_USE_RSYNC={use_rsync}
DOCKER_TIMEOUT={docker_timeout}
DOCKER_PULL_TIMEOUT={pull_timeout}
DOCKER_SAVE_TIMEOUT={save_timeout}
"""

    env_path = 'transfer_optimization.env'
    with open(env_path, 'w', encoding='utf-8') as f:
        f.write(env_text)

    print('\n已写入', env_path)
    print('请在应用已运行情况下，访问 历史 页面点“重载配置”，或调用:')
    print('  curl -X POST http://<host>:5000/api/reload_config')
    print('重载后新配置立即生效，无需重启。')

if __name__ == '__main__':
    main() 